/*
  This program implements a classical producer-consumer with a parallel queue.
 
  The additional lines of code for kOTF are appended with "*** kOTF ***"
 */


#include <iostream>
#include <assert.h>
#include <unistd.h>

#include <thread>
#include <mutex>
#include <condition_variable>
#include <atomic>
#include <string>
#include <utility>

#include <kotf/kotf.hpp>

using namespace kotf;
using namespace std;

State full_wait_state,
    empty_wait_state,
    push_state,
    pop_state,
    working_state;
Counter mutex_counter;
Channel full_chan,
    empty_chan;

unsigned producer_sleep = 100;
unsigned consumer_sleep = 100;


typedef unique_lock<mutex> Lock;


const int eof = -1;


class ParQueue {
    enum { max_size = 16 };
    
    int data[max_size];
    size_t start;
    size_t size;

    mutex dataMutex;
    condition_variable fullCondition, emptyCondition;

public:

    ParQueue()
    {
        reset();
    }

    void push(int s)
    {
        Stream& stream = Stream::mine();
        Scope local(stream, push_state); // delimits push()


        kotf::ScopedLock kotfGuard(stream, mutex_counter); // during the lock we see the counter "mutex_counter" with value=1
        Lock guard(dataMutex);
        kotfGuard.locked(); // now we locked the mutex, "mutex_counter" becomes 2.


        kotf::Condition queueFull(stream, full_wait_state, full_chan);
        while (size == max_size) {
            queueFull.wait();
            fullCondition.wait(guard);
        }
        queueFull.done();


        unsigned dst = (start + size) % max_size;
        data[dst] = s;
        if (size++ == 0) {
            stream.broadcast(empty_chan);
            emptyCondition.notify_all();
        }
    }

    int pop()
    {
        Stream& stream = Stream::mine();
        Scope local(stream, pop_state); // delimits pop()


        ScopedLock kotfGuard(stream, mutex_counter);
        Lock guard(dataMutex);
        kotfGuard.locked();

        
        Condition queueEmpty(stream, empty_wait_state, empty_chan);
        while (size == 0) {
            queueEmpty.wait();
            emptyCondition.wait(guard);
        }
        queueEmpty.done();


        int r = data[start];
        start = (start+1) % max_size;


        if (size-- == max_size) {
            stream.broadcast(full_chan);
            fullCondition.notify_all();
        }
        return r;
    }

    void reset()
    {
        start = 0;
        size = 0;
    }

    void stopConsumers()
    {
        Lock guard(dataMutex);
        while (size == max_size)
            fullCondition.wait(guard);

        unsigned dst = (start + size) % max_size;
        data[dst] = eof;
        if (size++ == 0)
            emptyCondition.notify_all();
    }
};

ParQueue pqueue;

atomic<unsigned> produced, consumed;


void producer(int start, int finish, int inc)
{
    string name = "producer-" + to_string(start);
    Stream::setName(name);
    Stream& stream = Stream::mine();

    int myprod = 0;
    {
        Scope working(stream, working_state);

        for (int i=start; i<finish; i += inc) {
            myprod += i;
            if (producer_sleep)
                usleep(producer_sleep); // sleep for a bit
            pqueue.push(i);
        }
    }
    produced += myprod;

    stream.finish();
}


void consumer(int id)
{
    string name = "consumer-" + to_string(id);
    Stream::setName(name);
    Stream& stream = Stream::mine();

    int mycons = 0;
    {
        Scope working(stream, working_state);

        while (true) {
            int s = pqueue.pop();
            if (s == eof) {
                pqueue.push(eof);
                break;
            }
            if (consumer_sleep)
                usleep(consumer_sleep); // sleep for a bit
            mycons += s;
        }
    }
    consumed += mycons;

    stream.finish();
}


void test(int n_producers, int n_consumers)
{
    // here we create all the kotf objects so we don't spend time later
    // creating them on the fly

    Session &session = Session::getDefault();

    push_state       = session.getState("push");
    pop_state        = session.getState("pop");
    full_wait_state  = session.getState("full");
    empty_wait_state = session.getState("empty");
    working_state    = session.getState("working");
    mutex_counter    = session.getCounter("mutex");
    full_chan        = session.getChannel("full");
    empty_chan       = session.getChannel("empty");


    produced = 0;
    consumed = 0;

    pqueue.reset();

    vector<thread> producers(n_producers), consumers(n_consumers);

    for (int i=0; i<n_producers; ++i)
        producers[i] = thread(producer, i, 1024, n_producers);
    
    for (int i=0; i<n_consumers; ++i)
        consumers[i] = thread(consumer, i);

    cout << "launched all threads" <<endl;

    for (thread& t : producers)
        t.join();
    cout << "joined with producers" << endl;

    pqueue.stopConsumers();

    cout << "pushed termination" << endl;
    
    for (thread& t : consumers)
        t.join();
    
    cout << "Produced: " << int(produced) << endl;
    cout << "Consumed: " << int(consumed) << endl;
    assert(int(produced) == int(consumed));
}


int main(int argc, char **argv)
{
    if (argc != 3) {
        cerr << "Usage: ./prod-cons prodsleep conssleep" << endl;
        cerr << "where\n"
            "prodsleep\t microseconds the producer will sleep\n"
            "conssleep\t microseconds the consumer will sleep\n" << endl;
        return 1;
    }

    producer_sleep = stoul(argv[1]);
    consumer_sleep = stoul(argv[2]);

    test(4, 4);
}
